#include "ray_ground_filter.hpp"
#include "gencolors.hpp"

// #include <pcl_ros/transforms.hpp>  // This header may not exist in some ROS versions
#include <pcl/filters/extract_indices.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl_conversions/pcl_conversions.h>

#include <cmath>
#include <algorithm>

namespace autoware_ground_segmentation
{

RayGroundFilter::RayGroundFilter(ros::NodeHandle& nh, ros::NodeHandle& private_nh)
: nh_(nh), private_nh_(private_nh), tf2_buffer_(ros::Duration(10)), tf2_listener_(tf2_buffer_)
{
  // Get parameters
  private_nh_.param<double>("general_max_slope", general_max_slope_, 8.0);
  private_nh_.param<double>("local_max_slope", local_max_slope_, 5.0);
  private_nh_.param<double>("initial_max_slope", initial_max_slope_, 3.0);
  private_nh_.param<double>("radial_divider_angle", radial_divider_angle_, 0.1);
  private_nh_.param<double>("concentric_divider_distance", concentric_divider_distance_, 0.01);
  private_nh_.param<double>("min_height_threshold", min_height_threshold_, 0.05);
  private_nh_.param<double>("reclass_distance_threshold", reclass_distance_threshold_, 0.1);
  private_nh_.param<double>("min_x", min_x_, -10.0);
  private_nh_.param<double>("max_x", max_x_, 10.0);
  private_nh_.param<double>("min_y", min_y_, -10.0);
  private_nh_.param<double>("max_y", max_y_, 10.0);
  private_nh_.param<bool>("use_vehicle_footprint", use_vehicle_footprint_, false);

  // Calculate derived parameters
  radial_dividers_num_ = static_cast<size_t>(360.0 / (radial_divider_angle_ * 180.0 / M_PI));
  
  // Setup vehicle footprint
  if (use_vehicle_footprint_) {
    setVehicleFootprint(min_x_, max_x_, min_y_, max_y_);
  }

  // Generate colors for visualization
  colors_ = generateColors(color_num_);

  // Setup PCL filters
  voxel_grid_.setLeafSize(0.1, 0.1, 0.1);

  // Setup ROS
  pointcloud_sub_ = nh_.subscribe("input", 1, &RayGroundFilter::pointcloudCallback, this);
  ground_cloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("output/ground", 1);
  no_ground_cloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("output/no_ground", 1);

  ROS_INFO("Ray Ground Filter initialized");
}

void RayGroundFilter::pointcloudCallback(const sensor_msgs::PointCloud2::ConstPtr& input_msg)
{
  // Convert to PCL
  pcl::PointCloud<PointType_>::Ptr input_cloud(new pcl::PointCloud<PointType_>);
  pcl::fromROSMsg(*input_msg, *input_cloud);

  if (input_cloud->empty()) {
    ROS_WARN("Empty input point cloud");
    return;
  }

  // Voxel grid filtering
  pcl::PointCloud<PointType_>::Ptr voxel_cloud(new pcl::PointCloud<PointType_>);
  voxel_grid_.setInputCloud(input_cloud);
  voxel_grid_.filter(*voxel_cloud);

  // Convert to organized format
  PointCloudXYZRTColor organized_points;
  std::vector<pcl::PointIndices> radial_divided_indices;
  std::vector<PointCloudXYZRTColor> radial_ordered_clouds;
  
  ConvertXYZIToRTZColor(voxel_cloud, organized_points, radial_divided_indices, radial_ordered_clouds);

  // Classify points
  pcl::PointIndices ground_indices, no_ground_indices;
  ClassifyPointCloud(radial_ordered_clouds, ground_indices, no_ground_indices);

  // Extract and publish results
  sensor_msgs::PointCloud2::Ptr ground_cloud_msg(new sensor_msgs::PointCloud2);
  sensor_msgs::PointCloud2::Ptr no_ground_cloud_msg(new sensor_msgs::PointCloud2);
  
  ExtractPointsIndices(input_msg, ground_indices, ground_cloud_msg, no_ground_cloud_msg);

  ground_cloud_pub_.publish(ground_cloud_msg);
  no_ground_cloud_pub_.publish(no_ground_cloud_msg);
}

void RayGroundFilter::ConvertXYZIToRTZColor(
  const pcl::PointCloud<PointType_>::Ptr in_cloud,
  PointCloudXYZRTColor & out_organized_points,
  std::vector<pcl::PointIndices> & out_radial_divided_indices,
  std::vector<PointCloudXYZRTColor> & out_radial_ordered_clouds)
{
  out_organized_points.resize(in_cloud->size());
  out_radial_divided_indices.clear();
  out_radial_divided_indices.resize(radial_dividers_num_);
  out_radial_ordered_clouds.clear();
  out_radial_ordered_clouds.resize(radial_dividers_num_);

  for (size_t i = 0; i < in_cloud->size(); ++i) {
    const auto& point = in_cloud->points[i];
    
    // Calculate radius and theta
    float radius = std::sqrt(point.x * point.x + point.y * point.y);
    float theta = std::atan2(point.y, point.x) * 180.0 / M_PI;
    if (theta < 0) theta += 360.0;
    
    // Determine radial division
    size_t radial_div = static_cast<size_t>(theta / (360.0 / radial_dividers_num_));
    if (radial_div >= radial_dividers_num_) radial_div = radial_dividers_num_ - 1;
    
    // Fill organized point
    PointXYZRTColor organized_point;
    organized_point.point = point;
    organized_point.radius = radius;
    organized_point.theta = theta;
    organized_point.radial_div = radial_div;
    organized_point.original_index = i;
    
    // Assign color
    size_t color_idx = radial_div % color_num_;
    organized_point.red = static_cast<size_t>(colors_[color_idx][2]);
    organized_point.green = static_cast<size_t>(colors_[color_idx][1]);
    organized_point.blue = static_cast<size_t>(colors_[color_idx][0]);
    
    out_organized_points[i] = organized_point;
    
    // Add to radial divisions
    out_radial_divided_indices[radial_div].indices.push_back(i);
    out_radial_ordered_clouds[radial_div].push_back(organized_point);
  }

  // Sort each radial division by radius
  for (auto& radial_cloud : out_radial_ordered_clouds) {
    std::sort(radial_cloud.begin(), radial_cloud.end(),
              [](const PointXYZRTColor& a, const PointXYZRTColor& b) {
                return a.radius < b.radius;
              });
  }
}

void RayGroundFilter::ClassifyPointCloud(
  std::vector<PointCloudXYZRTColor> & in_radial_ordered_clouds,
  pcl::PointIndices & out_ground_indices,
  pcl::PointIndices & out_no_ground_indices)
{
  out_ground_indices.indices.clear();
  out_no_ground_indices.indices.clear();

  for (auto& radial_cloud : in_radial_ordered_clouds) {
    if (radial_cloud.empty()) continue;

    // Initialize first point as ground
    if (!radial_cloud.empty()) {
      out_ground_indices.indices.push_back(radial_cloud[0].original_index);
    }

    for (size_t i = 1; i < radial_cloud.size(); ++i) {
      const auto& current_point = radial_cloud[i];
      const auto& previous_point = radial_cloud[i - 1];
      
      // Calculate height difference and distance
      float height_diff = current_point.point.z - previous_point.point.z;
      float distance = current_point.radius - previous_point.radius;
      
      if (distance < 0.001) continue; // Skip very close points
      
      // Calculate slope
      float slope = std::atan2(height_diff, distance) * 180.0 / M_PI;
      
      // Determine if point is ground based on slope thresholds
      bool is_ground = false;
      
      if (distance < reclass_distance_threshold_) {
        // Close points: use local slope threshold
        is_ground = std::abs(slope) < local_max_slope_;
      } else {
        // Far points: use general slope threshold
        is_ground = std::abs(slope) < general_max_slope_;
      }
      
      // Additional height threshold check
      if (std::abs(height_diff) < min_height_threshold_) {
        is_ground = true;
      }
      
      // Check vehicle footprint intersection if enabled
      if (use_vehicle_footprint_ && !is_ground) {
        Point point(current_point.point.x, current_point.point.y);
        auto intersection = calcPointVehicleIntersection(point);
        if (intersection) {
          // Point is inside vehicle footprint, classify as ground
          is_ground = true;
        }
      }
      
      if (is_ground) {
        out_ground_indices.indices.push_back(current_point.original_index);
      } else {
        out_no_ground_indices.indices.push_back(current_point.original_index);
      }
    }
  }
}

void RayGroundFilter::ExtractPointsIndices(
  const sensor_msgs::PointCloud2::ConstPtr in_cloud_ptr,
  const pcl::PointIndices & in_indices,
  sensor_msgs::PointCloud2::Ptr ground_cloud_msg_ptr,
  sensor_msgs::PointCloud2::Ptr no_ground_cloud_msg_ptr)
{
  // Initialize output clouds
  initializePointCloud2(in_cloud_ptr, ground_cloud_msg_ptr);
  initializePointCloud2(in_cloud_ptr, no_ground_cloud_msg_ptr);
  
  // Convert to PCL format for processing
  pcl::PointCloud<pcl::PointXYZ>::Ptr input_pcl(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::fromROSMsg(*in_cloud_ptr, *input_pcl);
  
  // Extract ground points
  pcl::PointCloud<pcl::PointXYZ>::Ptr ground_pcl(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::PointCloud<pcl::PointXYZ>::Ptr no_ground_pcl(new pcl::PointCloud<pcl::PointXYZ>);
  
  // Create temporary extract_indices objects
  pcl::ExtractIndices<pcl::PointXYZ> temp_extract;
  temp_extract.setInputCloud(input_pcl);
  temp_extract.setIndices(boost::make_shared<pcl::PointIndices>(in_indices));
  
  temp_extract.setNegative(false);
  temp_extract.filter(*ground_pcl);
  
  temp_extract.setNegative(true);
  temp_extract.filter(*no_ground_pcl);
  
  // Convert back to ROS messages
  pcl::toROSMsg(*ground_pcl, *ground_cloud_msg_ptr);
  pcl::toROSMsg(*no_ground_pcl, *no_ground_cloud_msg_ptr);
  
  ground_cloud_msg_ptr->header = in_cloud_ptr->header;
  no_ground_cloud_msg_ptr->header = in_cloud_ptr->header;
}

boost::optional<float> RayGroundFilter::calcPointVehicleIntersection(const Point & point)
{
  if (!use_vehicle_footprint_) {
    return boost::none;
  }
  
  // Check if point is inside vehicle footprint polygon
  if (boost::geometry::within(point, vehicle_footprint_)) {
    return 0.0f; // Point is inside, return distance 0
  }
  
  // Calculate distance to polygon boundary
  double distance = boost::geometry::distance(point, vehicle_footprint_);
  return static_cast<float>(distance);
}

void RayGroundFilter::setVehicleFootprint(const double min_x, const double max_x, const double min_y, const double max_y)
{
  vehicle_footprint_.clear();
  
  // Create rectangular footprint
  boost::geometry::append(vehicle_footprint_, Point(min_x, min_y));
  boost::geometry::append(vehicle_footprint_, Point(max_x, min_y));
  boost::geometry::append(vehicle_footprint_, Point(max_x, max_y));
  boost::geometry::append(vehicle_footprint_, Point(min_x, max_y));
  boost::geometry::append(vehicle_footprint_, Point(min_x, min_y)); // Close polygon
}

void RayGroundFilter::initializePointCloud2(
  const sensor_msgs::PointCloud2::ConstPtr & in_cloud_ptr,
  const sensor_msgs::PointCloud2::Ptr & out_cloud_msg_ptr)
{
  *out_cloud_msg_ptr = *in_cloud_ptr;
  out_cloud_msg_ptr->data.clear();
  out_cloud_msg_ptr->width = 0;
  out_cloud_msg_ptr->height = 0;
}

} // namespace autoware_ground_segmentation

int main(int argc, char** argv)
{
  ros::init(argc, argv, "ray_ground_filter_node");
  
  ros::NodeHandle nh;
  ros::NodeHandle private_nh("~");
  
  autoware_ground_segmentation::RayGroundFilter filter(nh, private_nh);
  
  ros::spin();
  
  return 0;
} 